#include <winsock2.h>
#include <mswsock.h>
#include <windows.h>
#include <iostream>
#include <thread>
#include <vector>
// #pragma comment(lib, "ws2_32.lib")
// #pragma comment(lib, "mswsock.lib")

#ifndef MAX_BUFFER_SIZE
#define MAX_BUFFER_SIZE 1024
#endif

struct PerIoData {
    OVERLAPPED overlapped;
    WSABUF wsabuf;
    char buffer[MAX_BUFFER_SIZE];
    SOCKET client_socket;
    bool is_accept;
    enum OperationType { NONE, RECV, SEND } op_type;
    PerIoData() : client_socket(INVALID_SOCKET), is_accept(false), op_type(NONE) {
        ZeroMemory(&overlapped, sizeof(OVERLAPPED));
        wsabuf.buf = buffer;
        wsabuf.len = MAX_BUFFER_SIZE;
    }
    ~PerIoData() {
        // Socket closing handled explicitly
    }
};

class IOCPServer {
public:
    IOCPServer() : iocp_(NULL), listen_socket_(INVALID_SOCKET) {
        WSADATA wsaData;
        WSAStartup(MAKEWORD(2, 2), &wsaData);

        // int num_threads = std::thread::hardware_concurrency(); // Get number of logical cores
        // if (num_threads == 0) num_threads = 1; // Fallback for systems that don't report cores

        // for (int i = 0; i < num_threads; ++i) {
        //     worker_threads_.emplace_back([this]() {
        //         this->Run(); // Each thread calls Run()
        //     });
        // }
    }

    ~IOCPServer() {
        closesocket(listen_socket_);
        if (iocp_) {
            CloseHandle(iocp_);
        }
        WSACleanup();
    }

    bool Start(const std::string& ip, unsigned short port) {
        // listen_socket_ = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        listen_socket_ = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
        if (listen_socket_ == INVALID_SOCKET) {
            std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
            return false;
        }

        sockaddr_in addr;
        addr.sin_family = AF_INET;
#define _WINSOCK_DEPRECATED_NO_WARNINGS
        // addr.sin_addr.s_addr = inet_addr(ip.c_str());
        addr.sin_addr.s_addr = htonl(INADDR_ANY); // 监听所有网卡

        addr.sin_port = htons(port);

        if (bind(listen_socket_, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {
            std::cerr << "Bind failed: " << WSAGetLastError() << std::endl;
            return false;
        }
        std::cout << "bind " << port << " succeeded.\n";

        if (listen(listen_socket_, SOMAXCONN) == SOCKET_ERROR) {
            std::cerr << "Listen failed: " << WSAGetLastError() << std::endl;
            return false;
        }
        std::cout << "listen succeeded.\n";

        iocp_ = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
        if (!iocp_) {
            std::cerr << "CreateIoCompletionPort failed: " << GetLastError() << std::endl;
            return false;
        }
        std::cout << "iocp_ in start:" << iocp_;


        CreateIoCompletionPort((HANDLE)listen_socket_, iocp_, 0, 0);
        for (int i = 0; i < 10; ++i) {
            PostAccept();
        }
        return true;
    }

    void Run() {
        std::cout << "Server started, waiting for events..." << std::endl;

        DWORD bytesTransferred;
        ULONG_PTR completionKey;
        LPOVERLAPPED overlapped;

        try {
            while (true) {
                std::cout << "1...\n";
                std::cout << "iocp_:" << iocp_;
                BOOL result = GetQueuedCompletionStatus(iocp_, &bytesTransferred, &completionKey, &overlapped, INFINITE);
                if (!result && !overlapped) {
                    DWORD err = GetLastError();
                    std::cerr << "GetQueuedCompletionStatus failed: " << err << std::endl;
                    continue;
                }

                std::cout << "2...\n";
                PerIoData* io_data = CONTAINING_RECORD(overlapped, PerIoData, overlapped);
                if (!result) {
                    DWORD err = WSAGetLastError();
                    std::cerr << "I/O operation failed: " << err << " for socket: " << io_data->client_socket << std::endl;
                    if (io_data->client_socket != INVALID_SOCKET && !io_data->is_accept) {
                        closesocket(io_data->client_socket);
                    }
                    if (io_data->is_accept) {
                        PostAccept();
                    }
                    delete io_data;
                    continue;
                }

                std::cout << "Completion event: is_accept=" << io_data->is_accept << ", bytes=" << bytesTransferred
                          << ", socket=" << io_data->client_socket << ", op_type=" << io_data->op_type << std::endl;

                if (io_data->is_accept) {
                    AcceptNewConnection(io_data, bytesTransferred);
                    PostAccept();
                } else {
                    if (io_data->op_type == PerIoData::RECV) {
                        HandleClientRequest(io_data, bytesTransferred);
                    } else if (io_data->op_type == PerIoData::SEND) {
                        std::cout << "WSASend completed for socket: " << io_data->client_socket << std::endl;
                        delete io_data;
                    }
                }
            }
        } catch (const std::exception& e) {
            std::cerr << "Exception in Run: " << e.what() << std::endl;
        } catch (...) {
            std::cerr << "Unknown exception in Run" << std::endl;
        }
    }

private:
    void PostAccept() {
        PerIoData* io_data = new PerIoData();
        io_data->is_accept = true;
        io_data->client_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        if (io_data->client_socket == INVALID_SOCKET) {
            std::cerr << "Client socket creation failed: " << WSAGetLastError() << std::endl;
            delete io_data;
            PostAccept(); // <--- ADD THIS: Immediately re-post if socket creation fails
            return;
        }

        LPFN_ACCEPTEX lpfnAcceptEx = nullptr;
        GUID GuidAcceptEx = WSAID_ACCEPTEX;
        DWORD bytes;
        if (WSAIoctl(listen_socket_, SIO_GET_EXTENSION_FUNCTION_POINTER,
                     &GuidAcceptEx, sizeof(GuidAcceptEx),
                     &lpfnAcceptEx, sizeof(lpfnAcceptEx),
                     &bytes, NULL, NULL) == SOCKET_ERROR) {
            std::cerr << "WSAIoctl failed: " << WSAGetLastError() << std::endl;
            closesocket(io_data->client_socket);
            delete io_data;
            PostAccept(); // <--- ADD THIS: Immediately re-post if WSAIoctl fails
            return;
        }

        char addr_buffer[(sizeof(sockaddr_in) + 16) * 2];
        ZeroMemory(addr_buffer, sizeof(addr_buffer));
        DWORD bytesReceived;
        BOOL result = lpfnAcceptEx(
            listen_socket_,
            io_data->client_socket,
            addr_buffer,
            0,
            sizeof(sockaddr_in) + 16,
            sizeof(sockaddr_in) + 16,
            &bytesReceived,
            &io_data->overlapped
        );

        if (!result && WSAGetLastError() != WSA_IO_PENDING) {
            std::cerr << "AcceptEx failed: " << WSAGetLastError() << std::endl;
            closesocket(io_data->client_socket);
            delete io_data;
            PostAccept(); // <--- ADD THIS: Crucial for maintaining accept pool
        } else {
            std::cout << "AcceptEx posted successfully for socket: " << io_data->client_socket << std::endl;
        }
    }

    void AcceptNewConnection(PerIoData* io_data, DWORD bytesTransferred) {
        std::cout << "New client connected! Socket: " << io_data->client_socket << std::endl;

        SOCKET client_socket = io_data->client_socket;
        CreateIoCompletionPort((HANDLE)client_socket, iocp_, (ULONG_PTR)client_socket, 0);

        if (setsockopt(client_socket, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT,
                       (char*)&listen_socket_, sizeof(listen_socket_)) == SOCKET_ERROR) {
            std::cerr << "setsockopt failed: " << WSAGetLastError() << std::endl;
            closesocket(client_socket);
            delete io_data;
            return;
        }

        PerIoData* recv_data = new PerIoData();
        recv_data->client_socket = client_socket;
        recv_data->is_accept = false;
        recv_data->op_type = PerIoData::RECV;
        io_data->client_socket = INVALID_SOCKET;

        DWORD flags = 0;
        int result = WSARecv(recv_data->client_socket, &recv_data->wsabuf, 1, NULL, &flags, &recv_data->overlapped, NULL);
        if (result == SOCKET_ERROR) {
            DWORD err = WSAGetLastError();
            if (err != WSA_IO_PENDING) {
                std::cerr << "WSARecv failed: " << err << " for socket: " << recv_data->client_socket << std::endl;
                closesocket(recv_data->client_socket);
                delete recv_data;
                delete io_data;
                return;
            }
        }
        std::cout << "WSARecv posted for socket: " << recv_data->client_socket << std::endl;

        delete io_data;
    }

    void HandleClientRequest(PerIoData* io_data, DWORD bytesTransferred) {
        SOCKET socket = io_data->client_socket;

        if (bytesTransferred == 0) {
            std::cout << "Client disconnected. Socket: " << socket << std::endl;
            closesocket(socket);
            delete io_data;
            return;
        }

        std::cout << "Received data: " << std::string(io_data->wsabuf.buf, bytesTransferred)
                  << " on socket: " << socket << std::endl;

        PerIoData* send_data = new PerIoData();
        send_data->client_socket = socket;
        send_data->is_accept = false;
        send_data->op_type = PerIoData::SEND;
        memcpy(send_data->wsabuf.buf, io_data->wsabuf.buf, bytesTransferred);
        send_data->wsabuf.len = bytesTransferred;

        DWORD bytesSent;
        int result = WSASend(socket, &send_data->wsabuf, 1, &bytesSent, 0, &send_data->overlapped, NULL);
        if (result == SOCKET_ERROR) {
            DWORD err = WSAGetLastError();
            if (err != WSA_IO_PENDING) {
                std::cerr << "WSASend failed: " << err << " for socket: " << socket << std::endl;
                closesocket(socket);
                delete send_data;
                delete io_data;
                return;
            }
        }
        std::cout << "WSASend posted for socket: " << socket << std::endl;

        // Always initiate new WSARecv to handle subsequent data
        PerIoData* new_io_data = new PerIoData();
        new_io_data->client_socket = socket;
        new_io_data->is_accept = false;
        new_io_data->op_type = PerIoData::RECV;

        DWORD flags = 0;
        result = WSARecv(new_io_data->client_socket, &new_io_data->wsabuf, 1, NULL, &flags, &new_io_data->overlapped, NULL);
        if (result == SOCKET_ERROR) {
            DWORD err = WSAGetLastError();
            if (err != WSA_IO_PENDING) {
                std::cerr << "WSARecv failed: " << err << " for socket: " << new_io_data->client_socket << std::endl;
                closesocket(new_io_data->client_socket);
                delete new_io_data;
                delete send_data;
                delete io_data;
                return;
            }
        }
        std::cout << "WSARecv posted for socket: " << new_io_data->client_socket << std::endl;

        delete io_data;
    }

    HANDLE iocp_;
    SOCKET listen_socket_;
    std::vector<std::thread> worker_threads_; // Store worker threads

};

int main() {
    try {
        IOCPServer server;
        if (!server.Start("0.0.0.0", 8080)) {
            std::cerr << "Failed to start the server!" << std::endl;
            return 1;
        }
        server.Run();
    } catch (const std::exception& e) {
        std::cerr << "Exception in main: " << e.what() << std::endl;
        return 1;
    } catch (...) {
        std::cerr << "Unknown exception in main" << std::endl;
        return 1;
    }
    return 0;
}